#!/usr/bin/python
# -*- coding: utf8 -*-

import sys
import os
import pprint

symbols = []
terminal_symbols = []
rules = []
terminal_rules = []

symbols_mapping = {}
terminal_symbols_mapping = {}
rules_mapping = {}
terminal_rules_mapping = {}

class Symbol:
    def __init__(self, name, terminal):
        self.name = name
        self.terminal = terminal
        if terminal:
            self.index = len(terminal_symbols)
            terminal_symbols.append(self)
            terminal_symbols_mapping[name] = self
        else:
            self.index = len(symbols)
            symbols.append(self)
            symbols_mapping[name] = self

    def string(self):
        return "%s%d_%s" % ('$' if self.terminal else '', self.index, self.name)

def get_symbol(query):
    if isinstance(query, int):
        return symbols[query]
    else:
        return symbols_mapping[query]

def get_terminal_symbol(query):
    if isinstance(query, int):
        return terminal_symbols[query]
    else:
        return terminal_symbols_mapping[query]


# represents S->XY or S->w
class Rule:
    def __init__(self, lhs, rhs):
        if len(rhs) == 1:
            self.index = len(terminal_rules)
            self.terminal = True
            self.s = get_symbol(lhs).index
            self.w = get_terminal_symbol(rhs[0]).index
            terminal_rules.append(self)
            terminal_rules_mapping[(self.s, self.w)] = self
        elif len(rhs) == 2:
            self.index = len(rules)
            self.terminal = False
            self.s = get_symbol(lhs).index
            self.x = get_symbol(rhs[0]).index
            self.y = get_symbol(rhs[1]).index
            rules.append(self)
            rules_mapping[(self.s, self.x, self.y)] = self
        else:
            raise
    
    def string(self):
        if self.terminal:
            return '%s->%s(#%d)' % (get_symbol(self.s).name,
                                    get_terminal_symbol(self.w).name,
                                    self.index)
        else:
            return '%s->%s,%s(#%d)' % (get_symbol(self.s).name,
                                       get_symbol(self.x).name,
                                       get_symbol(self.y).name,
                                       self.index)


def get_rule(query):
    if isinstance(query, int):
        return rules[query]
    else:
        return rules_mapping[query]

def get_terminal_rule(query):
    if isinstance(query, int):
        return terminal_rules[query]
    else:
        return terminal_rules_mapping[query]

def get_rules():
    return rules

terminal_rules_rlookup = None

def populate_terminal_rules_lookup():
    global terminal_rules_rlookup
    terminal_rules_rlookup = [[] for i in range(0, len(terminal_symbols))]
    for r in terminal_rules:
        terminal_rules_rlookup[r.w].append(r)

def get_terminal_rules_by_rhs(rhs):
    if isinstance(rhs, int):
        return terminal_rules_rlookup[rhs]
    else:
        return terminal_rules_rlookup[get_terminal_symbol(rhs).index]

def load_rules(filename):
    file = open(filename, 'r')
    lines = unicode(file.read(), 'gbk').split('\n')
    
    for line in lines:
        if line == '':
            continue

        main, at, prob = line.partition('@')
        fields = main.split()
        catagory = fields[0][1:]

        if catagory == 'TS':
            o = Symbol(fields[1], True)
        elif catagory == 'S':
            o = Symbol(fields[1], False)
        elif catagory == 'R':
            o = Rule(fields[1], fields[3:])
        else:
            raise
        if prob != None and prob != '':
            o.prob = float(prob)
        o.count = 0

    populate_terminal_rules_lookup()


if __name__ == '__main__':
    load_rules(sys.argv[1])
    for s in symbols:
        print s.string()
    for s in terminal_symbols:
        print s.string()
    for r in rules:
        print r.string() + ' ' + r.prob
    for r in terminal_rules:
        print r.string() + ' ' + r.prob

